package com.mrx.commons.trace;

import com.mrx.commons.utils.ThreadUtils;
import org.apache.commons.codec.binary.Hex;
import org.slf4j.MDC;
import org.springframework.cloud.commons.util.InetUtils;
import org.springframework.util.StringUtils;

import java.net.InetAddress;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author Mr.X
 * @since 2023-10-17 11:50
 **/
public class TraceUtils {

    private static final String TRACE_ID = "traceId";

    private static final String INDEX_NAME = "indexName";

    public static final String TRACE_HEADER = "x-trace";

    private static final AtomicInteger sequence = new AtomicInteger(1000);

    private static InetUtils inetUtils;

    public static void setInetUtils(InetUtils inetUtils) {
        TraceUtils.inetUtils = inetUtils;
    }

    public static void startTracer(String traceId) {
        startTracer(traceId, null);
    }

    public static void startTracer(String traceId, String indexName) {
        if (StringUtils.hasText(traceId)) {
            MDC.put(TRACE_ID, traceId);
            if (StringUtils.hasText(indexName)) {
                MDC.put(INDEX_NAME, indexName);
            }
        } else {
            startTracer(newTraceId(), indexName);
        }
    }

    public static void clearTracer() {
        MDC.remove(TRACE_ID);
    }

    public static Thread startTracer(Runnable r) {
        String traceId = getTraceId();
        return new Thread(() -> {
            try {
                startTracer(traceId);
                r.run();
            } finally {
                clearTracer();
            }
        });
    }

    public static ExecutorService newFixedTracerPool(int threads) {
        return newFixedTracerPool(threads, null);
    }

    public static ExecutorService newFixedTracerPool(int threads, String namePrefix) {
        return ThreadUtils.newFixedThreadPool(threads, TraceUtils::startTracer, namePrefix);
    }

    public static String getTraceId() {
        return Optional.ofNullable(MDC.get(TRACE_ID)).orElseGet(TraceUtils::newTraceId);
    }

    private static String newTraceId() {
        if (inetUtils != null) {
            return sofaTraceId();
        }
        String[] split = UUID.randomUUID().toString().split("-");
        return split[split.length - 1];
    }

    private static String sofaTraceId() {
        sequence.compareAndSet(9000, 1000);
        String ip = Optional.ofNullable(inetUtils.findFirstNonLoopbackAddress())
                .map(InetAddress::getAddress)
                .map(Hex::encodeHexString)
                .orElse(Hex.encodeHexString(InetAddress.getLoopbackAddress().getAddress()));
        long time = System.currentTimeMillis();
        long pid = ProcessHandle.current().pid();
        return ip + time + sequence.getAndIncrement() + pid;
    }

}
